#!/bin/bash

# Script to determine the memory footprint of the JVM.
# Run the script with root permissions in the same chroot environment where the
# team submissions are judged.
#
# Configure SANDBOX_MEM to the "memory limit" for the submissions (in MB).
# Set JVM_STACK to the stack size (-Xss) you want to allow for java submissions (in MB).
#
# Part of the DOMjudge Programming Contest Jury System and licensed
# under the GNU GPL. See README and COPYING for details.

cd "$(mktemp -d)" || exit -1

GROUPNAME=jvmtest
SUBGROUP=foo
JAVA_CLASS="Test"
SANDBOX_MEM=2048
JVM_STACK=64
JVM_HEAP=$((SANDBOX_MEM-JVM_STACK))
USED_HEAP=$((SANDBOX_MEM-JVM_STACK))
STEP=5

function write_and_compile_java() {
  FIRST_BLOCK=$((80*USED_HEAP/100))
  cat > "$JAVA_CLASS".java <<EOF
import java.util.*;
class Dummy {
  int value;
  Dummy next;
  Dummy(int value) {
    this.value = value;
    if (value % 2 == 0) {
      next = new Dummy(value - 1);
    }
  }
}
public class $JAVA_CLASS {
  public static final int USED_HEAP = $USED_HEAP;
  public static final int FIRST_BLOCK = $FIRST_BLOCK;
  public static int[][][] mem = new int[256][1024][FIRST_BLOCK];
  public static int[][][] mem2;
  public static void main(String[] args) throws Exception {
    try {
      ArrayList<Dummy> list = new ArrayList<>();
      for (int i = 0; i < 1000; i++) {
        int[][][] foo = new int[256][1024][5];
	for (int j = 0; j < 100; j++) {
	  list.add(new Dummy(i));
	  list.add(new Dummy(j));
	}
	foo[1][2][3] = list.get(42 + i).value;
      }
      list = null;
      mem2 = new int[256][1024][USED_HEAP - FIRST_BLOCK];
      stack();
    } catch (StackOverflowError soe) {
      System.err.println("\\\\o/");
    }
  }

  public static void stack() {
    stack();
  }
}
EOF
  javac "$JAVA_CLASS".java
}

function exec_java() {
  echo "- trying with: sandbox: ${SANDBOX_MEM}MB, stack: ${JVM_STACK}MB, heap: ${JVM_HEAP}MB, out of which ${USED_HEAP}MB will be used."
  cgcreate -g memory,cpu:"$GROUPNAME"
  cgcreate -g memory,cpu:"$GROUPNAME"/"$SUBGROUP"
  SANDBOX_BYTE=$((SANDBOX_MEM*1024*1024))
  echo "$SANDBOX_BYTE" > /sys/fs/cgroup/memory/"$GROUPNAME"/"$SUBGROUP"/memory.limit_in_bytes
  echo "$SANDBOX_BYTE" > /sys/fs/cgroup/memory/"$GROUPNAME"/"$SUBGROUP"/memory.memsw.limit_in_bytes
  cgexec -g memory,cpu:"$GROUPNAME"/"$SUBGROUP" java -XX:+UseSerialGC -Xmx"$JVM_HEAP"m -Xms"$JVM_HEAP"m -Xss"$JVM_STACK"m "$JAVA_CLASS" &> log
}

while true; do
	write_and_compile_java
	if exec_java; then
	  break;
	fi
	if grep OutOfMemoryError log; then
	  USED_HEAP=$((USED_HEAP-STEP))
	else
	  JVM_HEAP=$((JVM_HEAP-STEP))
	fi
	cgdelete -g memory,cpu:"$GROUPNAME" -r
done
MAX=$(cat /sys/fs/cgroup/memory/"$GROUPNAME"/"$SUBGROUP"/memory.memsw.max_usage_in_bytes)
echo "Memory usage in sandbox:" $((MAX/1024/1024)) "MB (should be close to ${SANDBOX_MEM}MB)"
echo "MEMRESERVED in java compile script should be at least" $((SANDBOX_MEM-JVM_HEAP)) "MB"
cgdelete -g memory,cpu:"$GROUPNAME" -r
